"""
@File    :  ssh_client.py
@Author  :  LeeSen
@Date    :  2024-06-30
@Version :  Python 3.7.4
"""
from paramiko.client import AutoAddPolicy, SSHClient
from paramiko.ssh_exception import AuthenticationException, NoValidConnectionsError


class SshClient:

    def __init__(self, host_ip, host_port, username, password, timeout):
        self.ssh_client = SSHClient()
        self.host_ip = host_ip
        self.username = username
        self.password = password
        self.port = host_port
        self.timeout = timeout

    def __enter__(self):
        try:
            self.ssh_client.set_missing_host_key_policy(AutoAddPolicy)

            self.ssh_client.connect(self.host_ip, self.port, self.username, self.password, timeout=self.timeout)

        except AuthenticationException as e:
            return 1001
        except NoValidConnectionsError as e:
            return 1002
        except Exception as a:
            return 1003
        return self

    def excute_command(self, commands):
        #TODO Exception
        try:
            stdin, stdout, stderr = self.ssh_client.exec_command(commands)
            return stdout.read().decode()
        except:
            return None

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.ssh_client.close()
